// 文件名：mdspan_gpu_pitch_adapter.cpp
// 目的：演示如何为“设备端按行对齐（pitch）分配的二维缓冲”构造 mdspan 视图，
//      并通过自定义 accessor 实现后端无关的只读访问（此处简化为主机模拟）。
// 组件：
//   - device_buffer：模拟设备内存分配，包含按 alignment 取整后的 pitch（以 float 元素计）；
//   - device_accessor：自定义 accessor，定义 data_handle_type/访问语义（这里保持简单直通）；
//   - device_mdspan：使用 layout_stride + device_accessor 的二维只读视图类型；
//   - print_device：按视图维度遍历打印元素。
// 关键点：
//   - strides = {pitch, 1} 指定行与列跨度（以元素为单位）；
//   - mdspan 的 mapping_type(layout_stride) 用 extents(rows, cols) + strides 生成索引映射；
//   - 设备端场景下，accessor 可封装设备加载、地址空间限定等（示例中为直通访问）。
#include <cstddef>
#include <iostream>
#include <mdspan>
#include <memory>

struct device_buffer {
    std::unique_ptr<float[], void(*)(float*)> ptr;
    std::size_t pitch; // 以 float 元素为单位的行跨度
    std::size_t rows;
    std::size_t cols;

    explicit device_buffer(std::size_t rows, std::size_t cols, std::size_t alignment = 256)
        : pitch(((cols * sizeof(float) + alignment - 1) / alignment) * alignment / sizeof(float)),
          rows(rows),
          cols(cols),
          ptr(static_cast<float*>(::operator new[](pitch * rows * sizeof(float), std::align_val_t(alignment))),
              [](float* p) { ::operator delete[](p, std::align_val_t(256)); }) {}
};

struct device_accessor {
    using element_type = const float;
    using reference = const float&;
    using data_handle_type = const float*;
    using offset_policy = device_accessor;

    constexpr device_accessor() noexcept = default;
    constexpr device_accessor offset(std::size_t) const noexcept { return *this; }
    constexpr reference access(data_handle_type p, std::size_t i) const noexcept { return p[i]; }
    constexpr data_handle_type decay(data_handle_type p) const noexcept { return p; }
};

using device_mdspan =
    std::mdspan<const float,
                std::extents<std::size_t, std::dynamic_extent, std::dynamic_extent>,
                std::layout_stride,
                device_accessor>;

void print_device(device_mdspan view) {
    for (std::size_t r = 0; r < view.extent(0); ++r) {
        for (std::size_t c = 0; c < view.extent(1); ++c) {
            std::cout << view(r, c) << ' ';
        }
        std::cout << '\n';
    }
}

int main() {
    device_buffer buf(3, 4);
    // 模拟设备写入：此处直接在主机指针上赋值
    for (std::size_t r = 0; r < buf.rows; ++r) {
        for (std::size_t c = 0; c < buf.cols; ++c) {
            buf.ptr.get()[r * buf.pitch + c] = static_cast<float>(r * 10 + c);
        }
    }

    std::array<std::size_t, 2> strides{buf.pitch, 1};
    device_mdspan view(buf.ptr.get(),
                       device_mdspan::mapping_type(
                           std::extents<std::size_t, std::dynamic_extent, std::dynamic_extent>(buf.rows, buf.cols),
                           strides));

    print_device(view);
    return 0;
}